#include "Base.h"
#include "BuiltinMeta.h"

#include <Detectron2/Utils/Utils.h>

using namespace std;
using namespace Detectron2;

///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

// All coco categories, together with their nice-looking visualization colors
// It's from https://github.com/cocodataset/panopticapi/blob/master/panoptic_coco_categories.json
struct Category {
	vector<float> color;
	int isthing;
	int id;
	string name;
};
static Category COCO_CATEGORIES[] = {
	{/* color */ { 220, 20, 60 }, /* isthing */ 1, /* id */ 1, /* name */ "person"},
	{/* color */ { 119, 11, 32 }, /* isthing */ 1, /* id */ 2, /* name */ "bicycle"},
	{/* color */ { 0, 0, 142 }, /* isthing */ 1, /* id */ 3, /* name */ "car"},
	{/* color */ { 0, 0, 230 }, /* isthing */ 1, /* id */ 4, /* name */ "motorcycle"},
	{/* color */ { 106, 0, 228 }, /* isthing */ 1, /* id */ 5, /* name */ "airplane"},
	{/* color */ { 0, 60, 100 }, /* isthing */ 1, /* id */ 6, /* name */ "bus"},
	{/* color */ { 0, 80, 100 }, /* isthing */ 1, /* id */ 7, /* name */ "train"},
	{/* color */ { 0, 0, 70 }, /* isthing */ 1, /* id */ 8, /* name */ "truck"},
	{/* color */ { 0, 0, 192 }, /* isthing */ 1, /* id */ 9, /* name */ "boat"},
	{/* color */ { 250, 170, 30 }, /* isthing */ 1, /* id */ 10, /* name */ "traffic light"},
	{/* color */ { 100, 170, 30 }, /* isthing */ 1, /* id */ 11, /* name */ "fire hydrant"},
	{/* color */ { 220, 220, 0 }, /* isthing */ 1, /* id */ 13, /* name */ "stop sign"},
	{/* color */ { 175, 116, 175 }, /* isthing */ 1, /* id */ 14, /* name */ "parking meter"},
	{/* color */ { 250, 0, 30 }, /* isthing */ 1, /* id */ 15, /* name */ "bench"},
	{/* color */ { 165, 42, 42 }, /* isthing */ 1, /* id */ 16, /* name */ "bird"},
	{/* color */ { 255, 77, 255 }, /* isthing */ 1, /* id */ 17, /* name */ "cat"},
	{/* color */ { 0, 226, 252 }, /* isthing */ 1, /* id */ 18, /* name */ "dog"},
	{/* color */ { 182, 182, 255 }, /* isthing */ 1, /* id */ 19, /* name */ "horse"},
	{/* color */ { 0, 82, 0 }, /* isthing */ 1, /* id */ 20, /* name */ "sheep"},
	{/* color */ { 120, 166, 157 }, /* isthing */ 1, /* id */ 21, /* name */ "cow"},
	{/* color */ { 110, 76, 0 }, /* isthing */ 1, /* id */ 22, /* name */ "elephant"},
	{/* color */ { 174, 57, 255 }, /* isthing */ 1, /* id */ 23, /* name */ "bear"},
	{/* color */ { 199, 100, 0 }, /* isthing */ 1, /* id */ 24, /* name */ "zebra"},
	{/* color */ { 72, 0, 118 }, /* isthing */ 1, /* id */ 25, /* name */ "giraffe"},
	{/* color */ { 255, 179, 240 }, /* isthing */ 1, /* id */ 27, /* name */ "backpack"},
	{/* color */ { 0, 125, 92 }, /* isthing */ 1, /* id */ 28, /* name */ "umbrella"},
	{/* color */ { 209, 0, 151 }, /* isthing */ 1, /* id */ 31, /* name */ "handbag"},
	{/* color */ { 188, 208, 182 }, /* isthing */ 1, /* id */ 32, /* name */ "tie"},
	{/* color */ { 0, 220, 176 }, /* isthing */ 1, /* id */ 33, /* name */ "suitcase"},
	{/* color */ { 255, 99, 164 }, /* isthing */ 1, /* id */ 34, /* name */ "frisbee"},
	{/* color */ { 92, 0, 73 }, /* isthing */ 1, /* id */ 35, /* name */ "skis"},
	{/* color */ { 133, 129, 255 }, /* isthing */ 1, /* id */ 36, /* name */ "snowboard"},
	{/* color */ { 78, 180, 255 }, /* isthing */ 1, /* id */ 37, /* name */ "sports ball"},
	{/* color */ { 0, 228, 0 }, /* isthing */ 1, /* id */ 38, /* name */ "kite"},
	{/* color */ { 174, 255, 243 }, /* isthing */ 1, /* id */ 39, /* name */ "baseball bat"},
	{/* color */ { 45, 89, 255 }, /* isthing */ 1, /* id */ 40, /* name */ "baseball glove"},
	{/* color */ { 134, 134, 103 }, /* isthing */ 1, /* id */ 41, /* name */ "skateboard"},
	{/* color */ { 145, 148, 174 }, /* isthing */ 1, /* id */ 42, /* name */ "surfboard"},
	{/* color */ { 255, 208, 186 }, /* isthing */ 1, /* id */ 43, /* name */ "tennis racket"},
	{/* color */ { 197, 226, 255 }, /* isthing */ 1, /* id */ 44, /* name */ "bottle"},
	{/* color */ { 171, 134, 1 }, /* isthing */ 1, /* id */ 46, /* name */ "wine glass"},
	{/* color */ { 109, 63, 54 }, /* isthing */ 1, /* id */ 47, /* name */ "cup"},
	{/* color */ { 207, 138, 255 }, /* isthing */ 1, /* id */ 48, /* name */ "fork"},
	{/* color */ { 151, 0, 95 }, /* isthing */ 1, /* id */ 49, /* name */ "knife"},
	{/* color */ { 9, 80, 61 }, /* isthing */ 1, /* id */ 50, /* name */ "spoon"},
	{/* color */ { 84, 105, 51 }, /* isthing */ 1, /* id */ 51, /* name */ "bowl"},
	{/* color */ { 74, 65, 105 }, /* isthing */ 1, /* id */ 52, /* name */ "banana"},
	{/* color */ { 166, 196, 102 }, /* isthing */ 1, /* id */ 53, /* name */ "apple"},
	{/* color */ { 208, 195, 210 }, /* isthing */ 1, /* id */ 54, /* name */ "sandwich"},
	{/* color */ { 255, 109, 65 }, /* isthing */ 1, /* id */ 55, /* name */ "orange"},
	{/* color */ { 0, 143, 149 }, /* isthing */ 1, /* id */ 56, /* name */ "broccoli"},
	{/* color */ { 179, 0, 194 }, /* isthing */ 1, /* id */ 57, /* name */ "carrot"},
	{/* color */ { 209, 99, 106 }, /* isthing */ 1, /* id */ 58, /* name */ "hot dog"},
	{/* color */ { 5, 121, 0 }, /* isthing */ 1, /* id */ 59, /* name */ "pizza"},
	{/* color */ { 227, 255, 205 }, /* isthing */ 1, /* id */ 60, /* name */ "donut"},
	{/* color */ { 147, 186, 208 }, /* isthing */ 1, /* id */ 61, /* name */ "cake"},
	{/* color */ { 153, 69, 1 }, /* isthing */ 1, /* id */ 62, /* name */ "chair"},
	{/* color */ { 3, 95, 161 }, /* isthing */ 1, /* id */ 63, /* name */ "couch"},
	{/* color */ { 163, 255, 0 }, /* isthing */ 1, /* id */ 64, /* name */ "potted plant"},
	{/* color */ { 119, 0, 170 }, /* isthing */ 1, /* id */ 65, /* name */ "bed"},
	{/* color */ { 0, 182, 199 }, /* isthing */ 1, /* id */ 67, /* name */ "dining table"},
	{/* color */ { 0, 165, 120 }, /* isthing */ 1, /* id */ 70, /* name */ "toilet"},
	{/* color */ { 183, 130, 88 }, /* isthing */ 1, /* id */ 72, /* name */ "tv"},
	{/* color */ { 95, 32, 0 }, /* isthing */ 1, /* id */ 73, /* name */ "laptop"},
	{/* color */ { 130, 114, 135 }, /* isthing */ 1, /* id */ 74, /* name */ "mouse"},
	{/* color */ { 110, 129, 133 }, /* isthing */ 1, /* id */ 75, /* name */ "remote"},
	{/* color */ { 166, 74, 118 }, /* isthing */ 1, /* id */ 76, /* name */ "keyboard"},
	{/* color */ { 219, 142, 185 }, /* isthing */ 1, /* id */ 77, /* name */ "cell phone"},
	{/* color */ { 79, 210, 114 }, /* isthing */ 1, /* id */ 78, /* name */ "microwave"},
	{/* color */ { 178, 90, 62 }, /* isthing */ 1, /* id */ 79, /* name */ "oven"},
	{/* color */ { 65, 70, 15 }, /* isthing */ 1, /* id */ 80, /* name */ "toaster"},
	{/* color */ { 127, 167, 115 }, /* isthing */ 1, /* id */ 81, /* name */ "sink"},
	{/* color */ { 59, 105, 106 }, /* isthing */ 1, /* id */ 82, /* name */ "refrigerator"},
	{/* color */ { 142, 108, 45 }, /* isthing */ 1, /* id */ 84, /* name */ "book"},
	{/* color */ { 196, 172, 0 }, /* isthing */ 1, /* id */ 85, /* name */ "clock"},
	{/* color */ { 95, 54, 80 }, /* isthing */ 1, /* id */ 86, /* name */ "vase"},
	{/* color */ { 128, 76, 255 }, /* isthing */ 1, /* id */ 87, /* name */ "scissors"},
	{/* color */ { 201, 57, 1 }, /* isthing */ 1, /* id */ 88, /* name */ "teddy bear"},
	{/* color */ { 246, 0, 122 }, /* isthing */ 1, /* id */ 89, /* name */ "hair drier"},
	{/* color */ { 191, 162, 208 }, /* isthing */ 1, /* id */ 90, /* name */ "toothbrush"},
	{/* color */ { 255, 255, 128 }, /* isthing */ 0, /* id */ 92, /* name */ "banner"},
	{/* color */ { 147, 211, 203 }, /* isthing */ 0, /* id */ 93, /* name */ "blanket"},
	{/* color */ { 150, 100, 100 }, /* isthing */ 0, /* id */ 95, /* name */ "bridge"},
	{/* color */ { 168, 171, 172 }, /* isthing */ 0, /* id */ 100, /* name */ "cardboard"},
	{/* color */ { 146, 112, 198 }, /* isthing */ 0, /* id */ 107, /* name */ "counter"},
	{/* color */ { 210, 170, 100 }, /* isthing */ 0, /* id */ 109, /* name */ "curtain"},
	{/* color */ { 92, 136, 89 }, /* isthing */ 0, /* id */ 112, /* name */ "door-stuff"},
	{/* color */ { 218, 88, 184 }, /* isthing */ 0, /* id */ 118, /* name */ "floor-wood"},
	{/* color */ { 241, 129, 0 }, /* isthing */ 0, /* id */ 119, /* name */ "flower"},
	{/* color */ { 217, 17, 255 }, /* isthing */ 0, /* id */ 122, /* name */ "fruit"},
	{/* color */ { 124, 74, 181 }, /* isthing */ 0, /* id */ 125, /* name */ "gravel"},
	{/* color */ { 70, 70, 70 }, /* isthing */ 0, /* id */ 128, /* name */ "house"},
	{/* color */ { 255, 228, 255 }, /* isthing */ 0, /* id */ 130, /* name */ "light"},
	{/* color */ { 154, 208, 0 }, /* isthing */ 0, /* id */ 133, /* name */ "mirror-stuff"},
	{/* color */ { 193, 0, 92 }, /* isthing */ 0, /* id */ 138, /* name */ "net"},
	{/* color */ { 76, 91, 113 }, /* isthing */ 0, /* id */ 141, /* name */ "pillow"},
	{/* color */ { 255, 180, 195 }, /* isthing */ 0, /* id */ 144, /* name */ "platform"},
	{/* color */ { 106, 154, 176 }, /* isthing */ 0, /* id */ 145, /* name */ "playingfield"},
	{/* color */ { 230, 150, 140 }, /* isthing */ 0, /* id */ 147, /* name */ "railroad"},
	{/* color */ { 60, 143, 255 }, /* isthing */ 0, /* id */ 148, /* name */ "river"},
	{/* color */ { 128, 64, 128 }, /* isthing */ 0, /* id */ 149, /* name */ "road"},
	{/* color */ { 92, 82, 55 }, /* isthing */ 0, /* id */ 151, /* name */ "roof"},
	{/* color */ { 254, 212, 124 }, /* isthing */ 0, /* id */ 154, /* name */ "sand"},
	{/* color */ { 73, 77, 174 }, /* isthing */ 0, /* id */ 155, /* name */ "sea"},
	{/* color */ { 255, 160, 98 }, /* isthing */ 0, /* id */ 156, /* name */ "shelf"},
	{/* color */ { 255, 255, 255 }, /* isthing */ 0, /* id */ 159, /* name */ "snow"},
	{/* color */ { 104, 84, 109 }, /* isthing */ 0, /* id */ 161, /* name */ "stairs"},
	{/* color */ { 169, 164, 131 }, /* isthing */ 0, /* id */ 166, /* name */ "tent"},
	{/* color */ { 225, 199, 255 }, /* isthing */ 0, /* id */ 168, /* name */ "towel"},
	{/* color */ { 137, 54, 74 }, /* isthing */ 0, /* id */ 171, /* name */ "wall-brick"},
	{/* color */ { 135, 158, 223 }, /* isthing */ 0, /* id */ 175, /* name */ "wall-stone"},
	{/* color */ { 7, 246, 231 }, /* isthing */ 0, /* id */ 176, /* name */ "wall-tile"},
	{/* color */ { 107, 255, 200 }, /* isthing */ 0, /* id */ 177, /* name */ "wall-wood"},
	{/* color */ { 58, 41, 149 }, /* isthing */ 0, /* id */ 178, /* name */ "water-other"},
	{/* color */ { 183, 121, 142 }, /* isthing */ 0, /* id */ 180, /* name */ "window-blind"},
	{/* color */ { 255, 73, 97 }, /* isthing */ 0, /* id */ 181, /* name */ "window-other"},
	{/* color */ { 107, 142, 35 }, /* isthing */ 0, /* id */ 184, /* name */ "tree-merged"},
	{/* color */ { 190, 153, 153 }, /* isthing */ 0, /* id */ 185, /* name */ "fence-merged"},
	{/* color */ { 146, 139, 141 }, /* isthing */ 0, /* id */ 186, /* name */ "ceiling-merged"},
	{/* color */ { 70, 130, 180 }, /* isthing */ 0, /* id */ 187, /* name */ "sky-other-merged"},
	{/* color */ { 134, 199, 156 }, /* isthing */ 0, /* id */ 188, /* name */ "cabinet-merged"},
	{/* color */ { 209, 226, 140 }, /* isthing */ 0, /* id */ 189, /* name */ "table-merged"},
	{/* color */ { 96, 36, 108 }, /* isthing */ 0, /* id */ 190, /* name */ "floor-other-merged"},
	{/* color */ { 96, 96, 96 }, /* isthing */ 0, /* id */ 191, /* name */ "pavement-merged"},
	{/* color */ { 64, 170, 64 }, /* isthing */ 0, /* id */ 192, /* name */ "mountain-merged"},
	{/* color */ { 152, 251, 152 }, /* isthing */ 0, /* id */ 193, /* name */ "grass-merged"},
	{/* color */ { 208, 229, 228 }, /* isthing */ 0, /* id */ 194, /* name */ "dirt-merged"},
	{/* color */ { 206, 186, 171 }, /* isthing */ 0, /* id */ 195, /* name */ "paper-merged"},
	{/* color */ { 152, 161, 64 }, /* isthing */ 0, /* id */ 196, /* name */ "food-other-merged"},
	{/* color */ { 116, 112, 0 }, /* isthing */ 0, /* id */ 197, /* name */ "building-other-merged"},
	{/* color */ { 0, 114, 143 }, /* isthing */ 0, /* id */ 198, /* name */ "rock-merged"},
	{/* color */ { 102, 102, 156 }, /* isthing */ 0, /* id */ 199, /* name */ "wall-other-merged"},
	{/* color */ { 250, 141, 255 }, /* isthing */ 0, /* id */ 200, /* name */ "rug-merged"},
};

void BuiltinMeta::_get_coco_instances_meta(Metadata &metadata) {
	vector<int> thing_ids;
	thing_ids.reserve(80);
	metadata->thing.reserve(80);
	for (auto k : COCO_CATEGORIES) {
		if (k.isthing == 1) {
			thing_ids.push_back(k.id);
			metadata->thing.push_back({ k.name, k.color });
		}
	}
	assert(thing_ids.size() == 80);

	// Mapping from the incontiguous COCO category id to an id in [0, 79]
	for (int i = 0; i < thing_ids.size(); i++) {
		metadata->thing_dataset_id_to_contiguous_id[thing_ids[i]] = i;
	}
}

void BuiltinMeta::_get_coco_panoptic_separated_meta(Metadata &metadata) {
	// 54 names for COCO stuff categories (including "things")
	metadata->stuff.reserve(54);
	// NOTE: I randomly picked a color for things
	metadata->stuff.push_back({ "things", { 82, 18, 128 } });

	vector<int> stuff_ids;
	stuff_ids.reserve(53);
	for (auto k : COCO_CATEGORIES) {
		if (k.isthing == 0) {
			stuff_ids.push_back(k.id);
			metadata->stuff.push_back({ replace_all(replace_all(k.name, "-other", ""), "-merged", ""), k.color });
		}
	}
	assert(stuff_ids.size() == 53);

    // For semantic segmentation, this mapping maps from contiguous stuff id
    // (in [0, 53], used in models) to ids in the dataset (used for processing results)
    // The id 0 is mapped to an extra category "thing".
	for (int i = 0; i < stuff_ids.size(); i++) {
		metadata->stuff_dataset_id_to_contiguous_id[stuff_ids[i]] = i + 1;
	}
    // When converting COCO panoptic annotations to semantic annotations
    // We label the "thing" category to 0
	metadata->stuff_dataset_id_to_contiguous_id[0] = 0;
}

void BuiltinMeta::_get_coco_person_meta(Metadata &metadata) {
	static vector<string> COCO_PERSON_KEYPOINT_NAMES = {
		"nose",
		"left_eye", "right_eye",
		"left_ear", "right_ear",
		"left_shoulder", "right_shoulder",
		"left_elbow", "right_elbow",
		"left_wrist", "right_wrist",
		"left_hip", "right_hip",
		"left_knee", "right_knee",
		"left_ankle", "right_ankle",
	};

	// Pairs of keypoints that should be exchanged under horizontal flipping
	static vector<tuple<string, string>> COCO_PERSON_KEYPOINT_FLIP_MAP = {
		{ "left_eye", "right_eye" },
		{ "left_ear", "right_ear" },
		{ "left_shoulder", "right_shoulder" },
		{ "left_elbow", "right_elbow" },
		{ "left_wrist", "right_wrist" },
		{ "left_hip", "right_hip" },
		{ "left_knee", "right_knee" },
		{ "left_ankle", "right_ankle" },
	};

	// rules for pairs of keypoints to draw a line between, and the line color to use.
	static vector<tuple<string, string, vector<float>>> KEYPOINT_CONNECTION_RULES = {
		// face
		{ "left_ear", "left_eye", { 102, 204, 255 }},
		{ "right_ear", "right_eye", { 51, 153, 255 }},
		{ "left_eye", "nose", { 102, 0, 204 }},
		{ "nose", "right_eye", { 51, 102, 255 }},
		// upper-body
		{ "left_shoulder", "right_shoulder", { 255, 128, 0 }},
		{ "left_shoulder", "left_elbow", { 153, 255, 204 }},
		{ "right_shoulder", "right_elbow", { 128, 229, 255 }},
		{ "left_elbow", "left_wrist", { 153, 255, 153 }},
		{ "right_elbow", "right_wrist", { 102, 255, 224 }},
		// lower-body
		{ "left_hip", "right_hip", { 255, 102, 0 }},
		{ "left_hip", "left_knee", { 255, 255, 77 }},
		{ "right_hip", "right_knee", { 153, 255, 204 }},
		{ "left_knee", "left_ankle", { 191, 255, 128 }},
		{ "right_knee", "right_ankle", { 255, 195, 77 }},
	};
	metadata->thing.push_back({ "person" });
	metadata->keypoint_names = COCO_PERSON_KEYPOINT_NAMES;
	metadata->keypoint_flip_map = COCO_PERSON_KEYPOINT_FLIP_MAP;
	metadata->keypoint_connection_rules = KEYPOINT_CONNECTION_RULES;
}

void BuiltinMeta::_get_cityscapes_meta(Metadata &metadata) {
	static vector<string> CITYSCAPES_THING_CLASSES = {
		"person", "rider", "car", "truck",
		"bus", "train", "motorcycle", "bicycle",
	};
	static vector<string> CITYSCAPES_STUFF_CLASSES = {
		"road", "sidewalk", "building", "wall", "fence", "pole", "traffic light",
		"traffic sign", "vegetation", "terrain", "sky", "person", "rider", "car",
		"truck", "bus", "train", "motorcycle", "bicycle", "license plate",
	};
	for (auto c : CITYSCAPES_THING_CLASSES) {
		metadata->thing.push_back({ c });
	}
	for (auto c : CITYSCAPES_STUFF_CLASSES) {
		metadata->stuff.push_back({ c });
	}
}

Metadata BuiltinMeta::_get_builtin_metadata(const std::string &dataset_name) {
	Metadata metadata(new MetadataImpl());
	if (dataset_name == "coco") {
		_get_coco_instances_meta(metadata);
	}
	else if (dataset_name == "coco_panoptic_separated") {
		_get_coco_instances_meta(metadata);
		_get_coco_panoptic_separated_meta(metadata);
	}
	else if (dataset_name == "coco_person") {
		_get_coco_person_meta(metadata);
	}
	else if (dataset_name == "cityscapes") {
		_get_cityscapes_meta(metadata);
	}
	else {
		assert(false); // "No built-in metadata for dataset {}".format(dataset_name))
	}
	return metadata;
}
